from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex



#读取文档，构造Node，用于生成检索评估数据集
documents = SimpleDirectoryReader(
                            input_files=["../../data/citys/南京市.txt"]).load_data()
node_parser = SentenceSplitter(chunk_size=1024)
nodes = node_parser.get_nodes_from_documents(documents)
for idx, node in enumerate(nodes):
    node.id_ = f"node_{idx}"

#准备一个检索器，后面使用
vector_index = VectorStoreIndex(nodes)
retriever = vector_index.as_retriever(similarity_top_k=2)

from llama_index.core.evaluation import (
    generate_question_context_pairs,
    EmbeddingQAFinetuneDataset,
)

QA_GENERATE_PROMPT_TMPL = """
以下是上下文:
---------------------
{context_str}
---------------------
你是一位专业教授。你的任务是基于以上的上下文，为即将到来的考试设置 {num_questions_per_chunk} 个问题。
这些问题必须基于提供的上下文生成，并确保上下文能够回答这些问题。确保每一行都只有一个独立的问题。不要有多余解释。不要给问题编号。"
"""

print("Generating question-context pairs...")
qa_dataset = generate_question_context_pairs(
    nodes, 
    llm=llm_ollama, 
    num_questions_per_chunk=1,
    qa_generate_prompt_tmpl=QA_GENERATE_PROMPT_TMPL
)    

print("Saving dataset...")
qa_dataset.save_json("retriever_eval_dataset.json")